from collections.abc import Iterator
from typing import Any, TypeVar, TypedDict

import numpy as np

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume

AnySample = TypeVar('AnySample')


class BufferedShuffledIteratorState(TypedDict):
    buffer: tuple[AnySample, ...]
    iterator_state: Any
    random_state: Any


class BufferedShuffledIterator(DatasetIterator[AnySample]):
    def __init__(self,
        iterator: Iterator[AnySample], *,
        buffer_size: int,
        random_seed: int = 42,
        state: BufferedShuffledIteratorState | None = None,
    ):
        self._iterator = iterator
        self._buffer_size = buffer_size
        if state is None:
            self._buffer = []
            self._random = np.random.default_rng(random_seed)
        else:
            self._buffer = list(state['buffer'])
            self._random = np.random.default_rng(random_seed)
            self._random.bit_generator.state = state['random_state']

    def __next__(self) -> AnySample:
        while len(self._buffer) < self._buffer_size:
            try:
                self._buffer.append(next(self._iterator))
            except StopIteration:
                break

        if len(self._buffer) == 0:
            raise StopIteration

        index = self._random.integers(len(self._buffer))
        return self._buffer.pop(index)

    def __pause__(self) -> BufferedShuffledIteratorState | None:
        if self._buffer is None:
            return None
        return BufferedShuffledIteratorState(
            buffer=tuple(self._buffer),
            iterator_state=pause(self._iterator),
            random_state=self._random.bit_generator.state)


class BufferedShuffledDataset(Dataset[AnySample]):
    def __init__(self,
        dataset: Dataset[AnySample], *,
        buffer_size: int = 1024,
        random_seed: int = 42,
    ):
        self._dataset = dataset
        self._buffer_size = buffer_size
        self._random_seed = random_seed

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def buffer_size(self) -> int:
        return self._buffer_size

    @property
    def random_seed(self) -> int:
        return self._random_seed

    def __len__(self) -> int:
        return len(self._dataset)

    def __iter__(self) -> DatasetIterator[AnySample]:
        return BufferedShuffledIterator(
            iterator=iter(self.dataset),
            buffer_size=self.buffer_size,
            random_seed=self.random_seed)

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[AnySample]:
        return BufferedShuffledIterator(
            iterator=resume(self.dataset, state['iterator_state']),
            buffer_size=self.buffer_size,
            random_seed=self.random_seed,
            state=state)
